from dataclasses import dataclass, field
from typing import Dict, List
from ..common.common import Simulator, CostCalculator

@dataclass
class BertBinaryHeadSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
        self.batch_size = mp['input_shape'][0][0]
        self.hidden_size = mp['input_shape'][0][1]
        self.h_out = mp['output_shape'][1]

    #linear没有切分，维度和输入X一致
    def get_calculate_workload(self, tp_size: int) -> List[int]:
        calculator = [0, 0, 0, 0, 0]
        # Linear 部分 总计算量=b×s×(h×hout​+hout​)
        calculator[0] += self.batch_size    # 加法计算次数
        calculator[2] += self.batch_size * self.hidden_size * self.h_out  # 乘法计算次数
        
        return calculator

    # 线性层的显存占用
    def get_memory_workload(self, tp_size: int) -> int:
        return 2 * self.hidden_size
    
    def get_communication_workload(self, tp_size) -> int:
        return 0
    
@dataclass
class PoolerSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
        self.hidden_size = mp['input_shape'][0][1]
        self.batch_size = mp['input_shape'][0][0]
        self.h_out = mp['output_shape'][1]

    def get_calculate_workload(self, tp_size: int) -> List[int]:
        calculator = [0, 0, 0, 0, 0]
        # Linear 部分 总计算量=b×s×(h×hout​+hout​)
        calculator[0] += self.batch_size   # 加法计算次数
        calculator[2] += self.batch_size * self.hidden_size * self.h_out  # 乘法计算次数
        
        return calculator

    def get_communication_workload(self, tp_size) -> int:
        return 0    
